/**
* \brief 
* \author pengcheng (pengcheng@yslrpch@126.com)
* \date 2020-05-30
* \attention CopyrightÃ‚Â©ADC Technology(tianjin)Co.Ltd
* \attention Refer to COPYRIGHT.txt for complete terms of copyright notice
*/

#ifndef DETECTION_VISION_CENTER_NET_DETECTOR_H__
#define DETECTION_VISION_CENTER_NET_DETECTOR_H__

#include <memory>
#include "detection_vision/impl/image_detector.hpp"
#include "detection_vision/impl/bbox2d.hpp"
#include <center_net.pb.h>
#include "detection_vision/tensorrt/tensorrt_inference.h"
#include "detection_vision/tensorrt/impl/utils.hpp"
namespace detection
{

template<typename Scalar>
class CenterNetDetector : public DetectorBase<Scalar>
{

public:
    using Ptr = std::shared_ptr<CenterNetDetector>;
    using ConstPtr = std::shared_ptr<const CenterNetDetector>;
    //using ImageBBox2dType = ImageBBox2D<Scalar>;
    typedef ImageBBox2D<Scalar> ImageBBox2dType;

    explicit CenterNetDetector(const std::string& config_file);
    virtual ~CenterNetDetector();

    virtual bool Detect(const std::vector<cv::Mat>&, std::vector<std::vector<ImageBBox2dType>>&);
private:
    void InitConfig(const std::string& config_file);

    void Inference(void);

    void Decode(std::vector<std::vector<ImageBBox2dType> >& objects);

    void PreProcess(const std::vector<cv::Mat>& images);

    std::string GetClassName(const unsigned int);

private:
    tensorrt_inference::CenterNetConfig config_;
    tensorrt_inference::TRTInfernce::Ptr p_trt_inference_;
    size_t batch_size_;
    size_t real_batch_size_;
    size_t cnn_input_width_;
    size_t cnn_input_height_;
    size_t cnn_input_channel_;
    size_t cnn_output_hm_channel_;
    size_t down_ratio_;
    size_t image_height_;
    size_t image_width_;
    float score_threshold_;
    float* input_buffer_; //only for one put header
    void* output_buffer_gpu_;
    char* output_buffer_cpu_;
    cudaStream_t stream_;
    float scale_ ;
    cv::Size sacle_size_;
    cv::Mat resize_mat_;
    cv::Mat crop_mat_;
    cv::Rect rect_crop_;
    cv::Mat float_mat_;
    size_t input_feature_map_size_;
    size_t input_batch_map_size_;
    std::vector<float> mean_;
    std::vector<float> stdd_;
};

}



#endif